# -*- coding: utf-8 -*-
# @Time    : 2023/5/8 15:32
# @Author  : Pan
# @Software: PyCharm
# @Project : VisualFramework
# @FileName: MAEDataset.py


import paddle
import numpy as np
from paddle import io
from datasets import base
from transforms import Compose


class MaskGenerator:
    def __init__(self, input_size=(224, 224), mask_patch_size=32, model_patch_size=4, mask_ratio=0.6):
        self.input_size = input_size
        self.mask_patch_size = mask_patch_size
        self.model_patch_size = model_patch_size
        self.mask_ratio = mask_ratio

        assert self.input_size[0] % self.mask_patch_size == 0 and self.input_size[1] % self.mask_patch_size == 0, "input_size 必须是 mask_patch_size 整数倍"
        assert self.mask_patch_size % self.model_patch_size == 0, "mask_patch_size 必须是 model_patch_size 整数倍"

        self.rand_size = (self.input_size[0] // self.mask_patch_size, self.input_size[1] // self.mask_patch_size)
        self.scale = self.mask_patch_size // self.model_patch_size

        self.token_count = (self.rand_size[0] * self.rand_size[1])
        self.mask_count = int(np.ceil(self.token_count * self.mask_ratio))

    def __call__(self):
        mask_idx = np.random.permutation(self.token_count)[:self.mask_count]
        mask = np.zeros(self.token_count, dtype=int)
        mask[mask_idx] = 1

        mask = mask.reshape(*self.rand_size)
        mask = mask.repeat(self.scale, axis=0).repeat(self.scale, axis=1)

        return mask


class MAEDataset(io.Dataset):
    def __init__(self, config):
        super(MAEDataset, self).__init__()
        self.mode = config["mode"] if "mode" in config.keys() else "standard"
        self.data_root = config["data_root"] if "data_root" in config.keys() else None
        self.recursion_identifier = config["recursion_identifier"] if "recursion_identifier" in config.keys() else None
        self.img_size = config["img_size"] if "img_size" in config.keys() else (224, 224)
        self.mask_patch_size = config["mask_patch_size"] if "mask_patch_size" in config.keys() else 32
        self.model_patch_size = config["model_patch_size"] if "model_patch_size" in config.keys() else 4
        self.mask_ratio = config["mask_ratio"] if "mask_ratio" in config.keys() else 0.6
        self.trans = Compose(config["transforms"])

        self.mask_generator = MaskGenerator(self.img_size, self.mask_patch_size, self.model_patch_size, self.mask_ratio)
        self.data_list = self._make_list(data_root=self.data_root)

    def __getitem__(self, item):
        data = eval("self._" + self.mode)(item)
        return data

    def __len__(self):
        return len(self.data_list)

    def _standard(self, item):
        data = {"img": self.data_list[item], "path": self.data_list[item]}
        data = self.trans(data)
        data["mask"] = self.mask_generator()
        return data

    def _predict(self, item):
        return {'img': paddle.randn((3, *self.img_size)), "path": "%08d.png" % item}

    def _make_list(self, data_root=None):
        data_list = base.recursion(data_root, self.recursion_identifier)
        return data_list
